package com.chenjt.six_rag;

import com.chenjt.three_tools.ModelUtil;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;
import org.springframework.boot.Banner;

import java.util.List;

/**
 * 注释：RAG检索数据
 *
 * @author chenjt
 * @date 2025/3/12 11:34
 */
public class MeituanRag {

    public static void main(String[] args) {

        String question = "忘记密码";// 用户提出问题

        //1.将问题进行向量化，并且去Redis中查询相似度比较高的知识条目
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .apiKey(ModelUtil.API_KEY)
                .baseUrl(ModelUtil.BASE_URL)
                .build();

        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536)// 这里因为embeddingModel维度是1536，所以这里写1536
                /**
                 * redis-cli命令行，以下命令需要安装stack
                 * FT._LIST：能看到索引列表
                 * FT.INFO meituan-rag：列出meituan-rag这个索引的详细信息
                 */
                .indexName("meituan-rag")
                .build();


        // 将embeddingModel和embeddingStore交给ContentRetriever管理
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .maxResults(5) // 最大匹配5个结果
                .minScore(0.8) // 匹配出来的内容要求相似度在0.8以上
                .build();

        // 查询相似度为0.8以上的条目
        Query query = Query.from(question);
        List<Content> contentList = contentRetriever.retrieve(query);
        for(Content content: contentList){
            System.out.println(content);
        }


        //2.将问题和相关知识条目整合成一个完整的信息，发给大模型
        DefaultContentInjector contentInjector = new DefaultContentInjector();
        /**
         * 通过提供的内容和用户的问题，通过模板生成完整信息，发送给大模型
         *  DefaultContentInjector默认的prompt模板（也就是System的提示词）
         *  DefaultContentInjector中的DEFAULT_PROMPT_TEMPLATE属性就是提供的模板，将提供的内容和用户的问题，注入到模板中，将内容提供给大模型
         */
        UserMessage promptMessage = contentInjector.inject(contentList, UserMessage.from(question));
        System.out.println(promptMessage.singleText());


        System.out.println("-----------------------------");

        ChatLanguageModel model = OpenAiChatModel.builder()
                .apiKey(ModelUtil.API_KEY)
                .baseUrl(ModelUtil.BASE_URL)
                .modelName("gpt-4o-mini")
                .build();
        Response<AiMessage> response = model.generate(promptMessage);
        System.out.println(response.content().text());


    }


}
